"""
任务状态管理器
"""
import json
import os
from datetime import datetime
from typing import Dict, Any, Optional
from enum import Enum
from core.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class StateManager:
    """状态管理器"""

    def __init__(self, pipeline_id: str, data_dir: str = "data/outputs"):
        self.pipeline_id = pipeline_id
        self.data_dir = data_dir
        self.state_file = os.path.join(data_dir, pipeline_id, "state.json")
        self.output_dir = os.path.join(data_dir, pipeline_id, "outputs")

        # 确保目录存在
        os.makedirs(self.output_dir, exist_ok=True)

        # 初始化状态
        self.state = {
            "pipeline_id": pipeline_id,
            "start_time": None,
            "end_time": None,
            "status": TaskStatus.PENDING.value,
            "nodes": {}
        }
        self._load_state()

    def _load_state(self):
        """加载状态文件"""
        if os.path.exists(self.state_file):
            try:
                with open(self.state_file, 'r', encoding='utf-8') as f:
                    self.state = json.load(f)
            except Exception as e:
                logger.warning(f"Failed to load state file: {e}")

    def _save_state(self):
        """保存状态文件"""
        try:
            with open(self.state_file, 'w', encoding='utf-8') as f:
                json.dump(self.state, f, indent=2, ensure_ascii=False, default=str)
        except Exception as e:
            logger.warning(f"Failed to save state file: {e}")

    def set_pipeline_status(self, status: TaskStatus):
        """设置整个pipeline的状态"""
        if self.state["start_time"] is None and status == TaskStatus.RUNNING:
            self.state["start_time"] = datetime.now().isoformat()
        elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
            self.state["end_time"] = datetime.now().isoformat()

        self.state["status"] = status.value
        self._save_state()

    def set_node_status(self, node_id: str, status: TaskStatus, result: Any = None, error: str = None):
        """设置节点状态"""
        if node_id not in self.state["nodes"]:
            self.state["nodes"][node_id] = {
                "status": TaskStatus.PENDING.value,
                "start_time": None,
                "end_time": None,
                "result_path": None,
                "error": None
            }

        node_state = self.state["nodes"][node_id]

        if node_state["start_time"] is None and status == TaskStatus.RUNNING:
            node_state["start_time"] = datetime.now().isoformat()
        elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
            node_state["end_time"] = datetime.now().isoformat()

        node_state["status"] = status.value

        if result is not None:
            # 保存结果到文件
            result_path = os.path.join(self.output_dir, f"{node_id}.json")
            try:
                with open(result_path, 'w', encoding='utf-8') as f:
                    json.dump(result, f, indent=2, ensure_ascii=False, default=str)
                node_state["result_path"] = result_path
            except Exception as e:
                logger.warning(f"Failed to save result for node {node_id}: {e}")

        if error is not None:
            node_state["error"] = str(error)

        self._save_state()

    def get_node_status(self, node_id: str) -> Optional[Dict[str, Any]]:
        """获取节点状态"""
        return self.state["nodes"].get(node_id)

    def get_pipeline_status(self) -> Dict[str, Any]:
        """获取整个pipeline的状态"""
        return self.state.copy()

    def get_node_result(self, node_id: str) -> Any:
        """获取节点执行结果"""
        node_state = self.state["nodes"].get(node_id)
        if not node_state or not node_state.get("result_path"):
            return None

        try:
            with open(node_state["result_path"], 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception:
            return None

    def get_completed_nodes(self) -> set:
        """获取已完成的节点集合"""
        completed = set()
        for node_id, node_state in self.state["nodes"].items():
            if node_state["status"] == TaskStatus.COMPLETED.value:
                completed.add(node_id)
        return completed